import taichi as ti

ti.init(arch=ti.cpu)

n=320
pixels = ti.field(dtype=ti.f32,shape=(n*2,n))

@ti.func
def complex_sqr(z):
    return ti.Vector([z[0]**2-z[1]**2,z[1]*z[0]*2])

@ti.kernel
def paint(t:ti.f32):
    for i,j in pixels:
        c = ti.Vector([-0.8,ti.cos(t)*0.2])
        z = ti.Vector([i/n-1,j/n-0.5])*2
        iter = 0
        while z.norm()<20 and iter<50:
            z=complex_sqr(z)+c
            iter+=1
        pixels[i,j]=1-iter*0.02

gui = ti.GUI("Julia Set",res=(n*2,n))

for i in range(1000000):
    paint(i*0.03)
    gui.set_image(pixels)
    gui.show()